#include <cuda_runtime_api.h> // cudaMalloc, cudaMemcpy, etc.
#include <cusparse.h>         // cusparseSpGEMM
#include <stdio.h>            // printf
#include <stdlib.h>           // EXIT_FAILURE = 1
                                                                              

#define CHECK_CUDA(func)                                                       \
{                                                                              \
    cudaError_t status = (func);                                               \
    if (status != cudaSuccess) {                                               \
        printf("CUDA API failed at line %d with error: %s (%d)\n",             \
               __LINE__, cudaGetErrorString(status), status);                  \
        return EXIT_FAILURE;                                                   \
    }                                                                          \
}



#define CHECK_CUSPARSE(func)                                                   \
{                                                                              \
    cusparseStatus_t status = (func);                                          \
    if (status != CUSPARSE_STATUS_SUCCESS) {                                   \
        printf("CUSPARSE API failed at line %d with error: %s (%d)\n",         \
               __LINE__, cusparseGetErrorString(status), status);              \
        return EXIT_FAILURE;                                                   \
    }                                                                          \
}                                                                              


//数组统一以 0 为起始下标
void printCsr(
    int m,
    int n,
    int nnz,
    const float *csrValA,
    const int *csrRowPtrA,
    const int *csrColIndA,
    const char* name)
{
    printf("matrix %s is %d-by-%d, nnz=%d, base=%d, \n", name, m, n, nnz, 0);
    for(int row = 0 ; row < m ; row++){
        printf("%d,%d\n",csrRowPtrA[row  ],csrRowPtrA[row+1]);
        const int start = csrRowPtrA[row  ];
        const int end   = csrRowPtrA[row+1];
        for(int colidx = start ; colidx < end ; colidx++){
            const int col = csrColIndA[colidx];
            const float Areg = csrValA[colidx];
            //printf("%s(%d,%d) = %f\n", name, row+1, col+1, Areg);
            printf("%s(%d,%d) = %f\n", name, row, col, Areg);
        }
    }                                                                                                                   
}


int cusparseMul(int A_num_rows, int A_num_cols, int *hA_csrOffsets, int *hA_columns, float *hA_values, int A_nnz,
		        int B_num_rows, int B_num_cols, int *hB_csrOffsets, int *hB_columns, float *hB_values, int B_nnz,
                int &C_num_rows, int &C_num_cols, int *hC_csrOffsets, int *hC_columns, float *hC_values, int &C_nnz)
{
	// 判断 A * B 是否满足相乘条件 或者是否有 0 矩阵
	if(A_num_cols != B_num_rows)
	{
		printf("%s\n","矩阵不满足相乘条件！");
		return EXIT_FAILURE;
	}
	if(A_nnz * B_nnz == 0)
	{
		printf("%s\n","有一个矩阵是 0 矩阵！");
		return EXIT_FAILURE;
	}

    // C = alpha * opA * opB + beta * C
    float               alpha       = 1.0f;
    float               beta        = 0.0f;                                                                             
    cusparseOperation_t opA         = CUSPARSE_OPERATION_NON_TRANSPOSE;
    cusparseOperation_t opB         = CUSPARSE_OPERATION_NON_TRANSPOSE;
    cudaDataType        computeType = CUDA_R_32F;//精度

    //--------------------------------------------------------------------------
    // GPU端分配内存：Device memory management: Allocate and copy A, B
    int   *dA_csrOffsets, *dA_columns, *dB_csrOffsets, *dB_columns,
          *dC_csrOffsets, *dC_columns;
    float *dA_values, *dB_values, *dC_values;


    // cudaError_t cudaMalloc(void** devPtr, size_t count);,返回分配内存的指针 *devPtr
    // allocate A
    CHECK_CUDA( cudaMalloc((void**) &dA_csrOffsets,
                           (A_num_rows + 1) * sizeof(int)) )
    CHECK_CUDA( cudaMalloc((void**) &dA_columns, A_nnz * sizeof(int))   )
    CHECK_CUDA( cudaMalloc((void**) &dA_values,  A_nnz * sizeof(float)) )
    // allocate B
    CHECK_CUDA( cudaMalloc((void**) &dB_csrOffsets,
                           (B_num_rows + 1) * sizeof(int)) )
    CHECK_CUDA( cudaMalloc((void**) &dB_columns, B_nnz * sizeof(int))   )
    CHECK_CUDA( cudaMalloc((void**) &dB_values,  B_nnz * sizeof(float)) )
	
    // allocate C offsets
    CHECK_CUDA( cudaMalloc((void**) &dC_csrOffsets,
                           (A_num_rows + 1) * sizeof(int)) )

    // CPU 数据 拷贝到 GPU
    // copy A
    CHECK_CUDA( cudaMemcpy(dA_csrOffsets, hA_csrOffsets,
                           (A_num_rows + 1) * sizeof(int),
                           cudaMemcpyHostToDevice) )
    CHECK_CUDA( cudaMemcpy(dA_columns, hA_columns, A_nnz * sizeof(int),
                           cudaMemcpyHostToDevice) )
    CHECK_CUDA( cudaMemcpy(dA_values, hA_values,
                           A_nnz * sizeof(float), cudaMemcpyHostToDevice) )
    // copy B
    CHECK_CUDA( cudaMemcpy(dB_csrOffsets, hB_csrOffsets,
                           (B_num_rows + 1) * sizeof(int),
                           cudaMemcpyHostToDevice) )
    CHECK_CUDA( cudaMemcpy(dB_columns, hB_columns, B_nnz * sizeof(int),
                           cudaMemcpyHostToDevice) )
    CHECK_CUDA( cudaMemcpy(dB_values, hB_values,
                           B_nnz * sizeof(float), cudaMemcpyHostToDevice) )


    //--------------------------------------------------------------------------
    // 调用cuSPARSE库，求解矩阵乘：CUSPARSE APIs
    // cuSparse 的上下文指针
    cusparseHandle_t     handle = NULL;
    // 稀疏矩阵描述符
    cusparseSpMatDescr_t matA, matB, matC;
    void*  dBuffer1    = NULL, *dBuffer2   = NULL;
    size_t bufferSize1 = 0,    bufferSize2 = 0;

    // 检查库 是否正确安装
    CHECK_CUSPARSE( cusparseCreate(&handle) )

    // Create sparse matrix A in CSR format
    // 以 CSR 格式初始化稀疏矩阵描述符matA
    CHECK_CUSPARSE( cusparseCreateCsr(&matA, A_num_rows, A_num_cols, A_nnz,
                                      dA_csrOffsets, dA_columns, dA_values,
                                      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
                                      CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F) )
    CHECK_CUSPARSE( cusparseCreateCsr(&matB, B_num_rows, B_num_cols, B_nnz,
                                      dB_csrOffsets, dB_columns, dB_values,
                                      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
                                      CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F) )
    CHECK_CUSPARSE( cusparseCreateCsr(&matC, A_num_rows, B_num_cols, 0,
                                      NULL, NULL, NULL,
                                      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
                                      CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F) )

    //--------------------------------------------------------------------------
    // SpGEMM Computation
    // 相当于spgemm的一个描述符  
    cusparseSpGEMMDescr_t spgemmDesc;
    CHECK_CUSPARSE( cusparseSpGEMM_createDescr(&spgemmDesc) )

    // cusparseSpGEMM_workEstimation()、cusparseSpGEMM_compute()用于确定缓冲区大小 和 执行实际的 计算
    // 请求外部存储器 bufferSize1 个字节：ask bufferSize1 bytes for external memory
    // 倒数第二个参数：workE…… 请求的工作空间字节数，最后一个参数 指向两个函数工作缓冲区的指针
    CHECK_CUSPARSE(
        cusparseSpGEMM_workEstimation(handle, opA, opB,
                                      &alpha, matA, matB, &beta, matC,
                                      computeType, CUSPARSE_SPGEMM_DEFAULT,
                                      spgemmDesc, &bufferSize1, NULL) )
    CHECK_CUDA( cudaMalloc((void**) &dBuffer1, bufferSize1) )

    // inspect the matrices A and B to understand the memory requirement for
    // the next step
    CHECK_CUSPARSE(
        cusparseSpGEMM_workEstimation(handle, opA, opB,
                                      &alpha, matA, matB, &beta, matC,
                                      computeType, CUSPARSE_SPGEMM_DEFAULT,
                                      spgemmDesc, &bufferSize1, dBuffer1) )


    // ask bufferSize2 bytes for external memory
    CHECK_CUSPARSE(
        cusparseSpGEMM_compute(handle, opA, opB,
                               &alpha, matA, matB, &beta, matC,
                               computeType, CUSPARSE_SPGEMM_DEFAULT,
                               spgemmDesc, &bufferSize2, NULL) )
    CHECK_CUDA( cudaMalloc((void**) &dBuffer2, bufferSize2) )

    // 计算A * B的中间过程 compute the intermediate product of A * B
    CHECK_CUSPARSE( cusparseSpGEMM_compute(handle, opA, opB,
                                           &alpha, matA, matB, &beta, matC,
                                           computeType, CUSPARSE_SPGEMM_DEFAULT,
                                           spgemmDesc, &bufferSize2, dBuffer2) )

    //到此步 已经计算完成！

    // get matrix C non-zero entries C_nnz1
    // 通过 matC 描述符，得到其余三个 值,行数、列数、非 0 个数
    int64_t C_num_rows1, C_num_cols1, C_nnz1;
    CHECK_CUSPARSE( cusparseSpMatGetSize(matC, &C_num_rows1, &C_num_cols1,
                                         &C_nnz1) )
    // allocate matrix C
    CHECK_CUDA( cudaMalloc((void**) &dC_columns, C_nnz1 * sizeof(int))   )
    CHECK_CUDA( cudaMalloc((void**) &dC_values,  C_nnz1 * sizeof(float)) )

	// update matC with the new pointers
    // 更新 matC,通过后三个参数
    CHECK_CUSPARSE(
        cusparseCsrSetPointers(matC, dC_csrOffsets, dC_columns, dC_values) )

    // if beta != 0, cusparseSpGEMM_copy reuses/updates the values of dC_values
    // 复制最终的结果给矩阵C：copy the final products to the matrix C
    // 最终结果 存在了 matC 的 三个数组里 ？
    CHECK_CUSPARSE(
        cusparseSpGEMM_copy(handle, opA, opB,
                            &alpha, matA, matB, &beta, matC,
                            computeType, CUSPARSE_SPGEMM_DEFAULT, spgemmDesc) )

    // destroy matrix/vector descriptors
    CHECK_CUSPARSE( cusparseSpGEMM_destroyDescr(spgemmDesc) )
    CHECK_CUSPARSE( cusparseDestroySpMat(matA) )
    CHECK_CUSPARSE( cusparseDestroySpMat(matB) )
    CHECK_CUSPARSE( cusparseDestroySpMat(matC) )
    CHECK_CUSPARSE( cusparseDestroy(handle) )

	//--------------------------------------------------------------------------
    // 检查结果是否计算正确： device result check
    // int64_t 转化为 int 
	C_num_rows = C_num_rows1;
	C_num_cols = C_num_cols1;
    C_nnz      = C_nnz1;
	
	hC_csrOffsets = new int[A_num_rows + 1];
    hC_columns    = new int[C_nnz];
    hC_values     = new float[C_nnz];
    CHECK_CUDA( cudaMemcpy(hC_csrOffsets, dC_csrOffsets,
                           (A_num_rows + 1) * sizeof(int),
                           cudaMemcpyDeviceToHost) )
    CHECK_CUDA( cudaMemcpy(hC_columns, dC_columns, C_nnz * sizeof(int),
                           cudaMemcpyDeviceToHost) )
    CHECK_CUDA( cudaMemcpy(hC_values, dC_values, C_nnz * sizeof(float),
                           cudaMemcpyDeviceToHost) )
    
	printCsr(C_num_rows, C_num_cols, C_nnz, hC_values, hC_csrOffsets, hC_columns, "C");
/*
	for(int i = 0; i < C_nnz ; i++)
    {
         printf("%d\n",hC_columns_tmp[i]);
    }
    printf("%s\n","-------------");

    for(int i = 0; i < C_nnz ; i++)
    {
         printf("%f\n",hC_values_tmp[i]);
    }
*/
    //--------------------------------------------------------------------------
    // device memory deallocation
    CHECK_CUDA( cudaFree(dBuffer1) )
    CHECK_CUDA( cudaFree(dBuffer2) )
    CHECK_CUDA( cudaFree(dA_csrOffsets) )
    CHECK_CUDA( cudaFree(dA_columns) )
    CHECK_CUDA( cudaFree(dA_values) )
    CHECK_CUDA( cudaFree(dB_csrOffsets) )
    CHECK_CUDA( cudaFree(dB_columns) )
    CHECK_CUDA( cudaFree(dB_values) )
    CHECK_CUDA( cudaFree(dC_csrOffsets) )
    CHECK_CUDA( cudaFree(dC_columns) )
    CHECK_CUDA( cudaFree(dC_values) )

	return 0;
}

int main()
{

    int A_num_rows = 4;
    int A_num_cols = 4;
    int A_nnz      = 9;
    int B_num_rows = 4;
    int B_num_cols = 4;
    int B_nnz      = 9;
    int   hA_csrOffsets[] = { 0, 3, 4, 7, 9 };
    int   hA_columns[]    = { 0, 2, 3, 1, 0, 2, 3, 1, 3 };
    float hA_values[]     = { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f,
                              6.0f, 7.0f, 8.0f, 9.0f };
    int   hB_csrOffsets[] = { 0, 2, 4, 7, 8 };
    int   hB_columns[]    = { 0, 3, 1, 3, 0, 1, 2, 1 };
    float hB_values[]     = { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f,
                              6.0f, 7.0f, 8.0f };
	
	int   C_num_rows;
	int   C_num_cols;
	int   C_nnz;
	int   *hC_csrOffsets = NULL;
    int   *hC_columns    = NULL;
    float *hC_values     = NULL;
     
	int status = cusparseMul(A_num_rows, A_num_cols, hA_csrOffsets, hA_columns, hA_values, A_nnz,
                             B_num_rows, B_num_cols, hB_csrOffsets, hB_columns, hB_values, B_nnz,
                             C_num_rows, C_num_cols, hC_csrOffsets, hC_columns, hC_values, C_nnz);
	if(status)
	{
		printf("%s,\n","程序未成功执行！");
		return 0;
	}
    //printf("%d",C_nnz);
  //  printCsr(C_num_rows, C_num_cols, C_nnz, hC_values, hC_csrOffsets, hC_columns, "C"); 
/*	
	for(int i = 0; i < C_nnz ; i++)
    {
         printf("%d\n",hC_columns[i]);
    }
    printf("%s\n","-------------");

    for(int i = 0; i < C_nnz ; i++)
    {
         printf("%f\n",hC_values[i]);
    }

	*/
	return 0;

}
